import sys
sys.path.append('/home/FAST_DATA_MIRROR/Langchain-Chatchat-master')
from tabel_and_images.read_image import parse_image_llava
from PIL import Image, ImageDraw
import fitz  # PyMuPDF
import os
from PIL import Image, ImageDraw
from pathlib import Path
import shutil
from collections import Counter
from langchain.document_loaders.unstructured import UnstructuredFileLoader
from typing import List
from server.utils import wrap_done, get_ChatOpenAI
from langchain.chat_models import ChatOpenAI
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import HumanMessage, AIMessage
chat_model : BaseChatModel = get_ChatOpenAI(
    model_name='chatglm3-6b',
    temperature=0.7,
    max_tokens=None,
)

def calculate_average_height(page):
    # 获取页面上所有文本块及其位置
    text_blocks = page.get_text("blocks")
    # 计算所有文本块的高度
    heights = [block[3] - block[1] for block in text_blocks]
    # 统计最常见的高度
    height_counter = Counter(heights)
    # 找出出现最频繁的高度
    if height_counter:
        most_common_height = height_counter.most_common(1)[0][0]
        return most_common_height
    else:
        return 0
def is_rect_inside(rect_a, rect_b):
    """
    判断矩形A(rect_a)是否完全在矩形B(rect_b)内部。
    矩形用四元组表示：(x0, y0, x1, y1)，其中(x0, y0)为左上角坐标，(x1, y1)为右下角坐标。
    """
    return rect_a[0] >= rect_b[0] and rect_a[1] >= rect_b[1] and rect_a[2] <= rect_b[2] and rect_a[3] <= rect_b[3]
def find_captions(page, images_rects):
    captions = []
    average_height = calculate_average_height(page)
    search_gap = average_height * 2  # 假设两行文字的平均高度
    # 使用get_text('blocks')获取文本块及其位置
    text_blocks = page.get_text("blocks")
    for img_rect in images_rects:
        bottom = img_rect[3]  # 图片底边的y坐标
        # 初始化变量用于确定与图片最接近的题注
        closest_caption = None
        shortest_distance = float('inf')
        for block in text_blocks:
            block_rect = block[:4]  # 文本块的矩形位置包含在block的前四个元素
            text = block[4]  # 文本内容
            # 检查文本块是否位于图片下方的预定搜索区间内
            if block_rect[1] > bottom and block_rect[1] < bottom + search_gap:
                # 检查文本是否以“图”或“Figure”开头
                if text.lower().startswith(('图', 'figure')):
                    # 计算文本块顶部与图片底部的距离
                    distance = block_rect[1] - bottom
                    # 检查文本块宽度与图片宽度的接近程度
                    img_width = img_rect[2] - img_rect[0]
                    caption_width = block_rect[2] - block_rect[0]
                    width_difference = abs(img_width - caption_width)
                    # 选择距离最短且宽度差最小的题注
                    if distance < shortest_distance or (distance == shortest_distance and width_difference < shortest_gap):
                        closest_caption = block
                        shortest_distance = distance
                        shortest_gap = width_difference
        if closest_caption and closest_caption not in captions:
            captions.append(closest_caption)
    return captions
def highlight_rectangles(page, draw=None):
    all_rects = []
    # 首先获取所有矩形
    for drawing in page.get_drawings():
        if drawing.get("type") == "f" and drawing.get("rect"):
            bbox = drawing["rect"]
            all_rects.append(bbox)
    # 过滤掉被其他矩形包含的矩形
    rects_to_draw = []
    for rect in all_rects:
        if not any(is_rect_inside(rect, other_rect) for other_rect in all_rects if rect != other_rect):
            rects_to_draw.append(rect)
    captions = find_captions(page, rects_to_draw)
    # 画出最外层的矩形
    if draw:
        for bbox in rects_to_draw:
            draw.rectangle([bbox[0], bbox[1], bbox[2], bbox[3]], outline="blue")
        for i, cap in enumerate(captions):
            draw.rectangle([cap[0], cap[1], cap[2], cap[3]], outline="red")  # 使用红色标出题注位置        
    return rects_to_draw, captions

def extract_text(pdf_file):
    tard = Path(pdf_file).with_suffix('')
    tard.mkdir(exist_ok=True)
    doc = fitz.open(pdf_file)
    body_text = []
    page_text = {}
    page_images_all = {}
    for page_num in range(len(doc)):
        body_text_page = []
        page_images_all[page_num] = []
        page = doc.load_page(page_num)
        pix = doc.get_page_pixmap(page_num)
        mode = "RGBA" if pix.alpha else "RGB"    
        img_big = Image.frombytes(mode, [pix.width, pix.height], pix.samples)
        text_blocks = page.get_text("blocks")
        image_rects, captions = highlight_rectangles(page)
        for image_i, (box, cap) in enumerate(zip(image_rects, captions)):
            tar = tard.joinpath(f'{page_num}-{image_i}.png').__str__()
            img_big.crop(box).save(tar)
            page_images_all[page_num].append([tar, cap[4]])
        exclusion_rects = image_rects + [cap[:4] for cap in captions]
        for block in text_blocks:
            rect, text, font_size = block[:4], block[4], block[5]
            if any(is_rect_inside(rect, exc_rect) for exc_rect in exclusion_rects):
                continue
            body_text.append(text)
            body_text_page.append(text)
        page_text[page_num] = body_text_page
    tabel_contexts = []
    for page_num in range(len(doc)):
        context_page_pre, context_page_next = [''], ['']
        if page_num>0:
            text_list = page_text[page_num-1]
            context_page_pre = text_list[len(text_list)//2:]
        if page_num<len(doc)-1:
            text_list = page_text[page_num+1]
            context_page_next = text_list[:len(text_list)//2]
        context_page_now = '\n'.join(context_page_pre+page_text[page_num]+context_page_next)
        context_page_now = context_page_now.replace('-\n','')
        context_page_now = context_page_now.replace('\n','')            
        # 制作表格数据
        promt = f'''
            文章内容: {context_page_now},
            请根据文章内容判断这段文章中是否有表格数据,如果有表格,则提取出表格,如果没有表格,则返回空。
            表格内容:
        '''
        messages = [
            [HumanMessage(content=promt)],
        ]
        out = chat_model.generate(messages=messages)
        out = out.generations[0][0].text
        tabel_contexts.append(out)
        # 制作表格数据
        for image, cap_text in page_images_all[page_num]:
            promt = f'''
                图片题注: {cap_text},
                图片上下文: {context_page_now},
                请根据图片题注和图片上下文,生成一个符合逻辑和上下文语境的图片描述。
                图片描述:
            '''
            text_llava = parse_image_llava(image, message=promt)
            text_llava = '\n' + '-'*100 + '\n' + cap_text + ':' + text_llava
            body_text.append(text_llava)
    body_text += tabel_contexts
    body_text = ''.join(body_text)
    return body_text

class Pdf_Image_Loader(UnstructuredFileLoader):
    def _get_elements(self) -> List:
        text = extract_text(self.file_path)
        from unstructured.partition.text import partition_text
        return partition_text(text=text, **self.unstructured_kwargs)


if __name__=='__main__':
    pdf = 'tabel_and_images/1904.02701v1.pdf'
    body = extract_text(pdf)
    with open('body.txt', 'w', encoding='utf-8') as f1:
        for paragraph in body:
            f1.write(paragraph+'\n')